import os
import cv2
import matplotlib.pyplot as plt
from matplotlib import rcParams
from collections import Counter
from ultralytics import YOLO

# 设置支持中文的字体
rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['SimHei']  # 或者使用其他支持中文的字体

def infer(image_path, model):
    # 加载并处理输入图像
    img = cv2.imread(image_path)
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # 将图像转换为RGB模式

    # 执行推理
    results = model(img_rgb)

    # 统计每个类别的数量
    class_counts = Counter()

    # 设置置信度阈值
    conf_threshold = 0.75

    # 输出检测的类别、置信度和边界框信息
    for detection in results[0].boxes.data:
        x1, y1, x2, y2, score, class_id = detection.tolist()

        # 仅处理置信度大于或等于 0.8 的检测
        if score >= conf_threshold:
            class_name = model.names[int(class_id)]
            class_counts[class_name] += 1

    # 打印每个类别及其数量，不换行
    for class_name, count in class_counts.items():
        print(f"{class_name}:{count}", end='  ')
    print()  # 打印完后加一个空行

    # 可视化检测结果，滤掉低置信度结果
    result_img = results[0].plot()  # 绘制检测结果
    plt.imshow(result_img)
    plt.title(f'推理结果: {os.path.basename(image_path)}')
    plt.axis('off')
    plt.show()

def main():
    # 文件夹路径
    folder_path = r'captured_images'  # 替换为包含图片的文件夹路径
    image_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(('.jpg', '.png', '.jpeg'))]

    # 加载训练好的YOLO模型
    model = YOLO('./ultralytics-main/best.pt')  # 使用训练后保存的模型

    # 遍历文件夹中的所有图片文件
    for image_path in image_files:
        infer(image_path, model)

if __name__ == '__main__':
    main()
